/* -*- c++ -*- */

#ifndef _BIT_ARRAY_H_
#define _BIT_ARRAY_H_

#include "defs.h"

#ifdef _OPENMP
#include <omp.h>
#else // _OPENMP
#include "no_omp.h"
#endif // _OPENMP

// #include <assert.h>

class BitArray {
 public:
  BitArray( _key_t tam): _size( 0), _lst_bits( NULL) {
    _size = ( size_t)( ( tam + 31 ) >> 5);
    _lst_bits = new unsigned int[ _size];
  }
  BitArray( const BitArray &ba): _size( 0), _lst_bits( NULL) {
    _size = ba._size;
    if ( _size) {
      _lst_bits = new unsigned int[ _size];
      for ( size_t i = 0; i < _size; i++) {
        _lst_bits[ i] = ba._lst_bits[ i];
      }
    }
  }
  // BitArray( long long int tam): _size( 0), _lst_bits( NULL) {
  //   _size = ( ( tam + 63 ) >> 6) & 0xffffffffffffffffL;
  //   _lst_bits = new unsigned int[ _size];
  // }
  ~BitArray() {
    delete[] _lst_bits;
  }

  BitArray &operator=( const BitArray &ba) {
    delete[] _lst_bits;
    _lst_bits = NULL;
    _size = ba._size;
    if ( _size) {
      _lst_bits = new unsigned int[ _size];
      for ( size_t i = 0; i < _size; i++) {
        _lst_bits[ i] = ba._lst_bits[ i];
      }
    }
    return *this;
  }

  _key_t size() const {
    return ( _key_t)_size << 5; // it's not exact, but...
  }
  
  unsigned int Get( _key_t idx) const {
    size_t g_idx = idx >> 5;
    int l_idx = idx & 0x1f;
    // assert( g_idx < _size);
    return ( _lst_bits[ g_idx] & ( 1 << l_idx)) >> l_idx;
  }

  // long long int Get( long long int idx) const {
  //   size_t g_idx = ( idx >> 6) & 0xffffffffffffffffL;
  //   long long int l_idx = idx & 0x3f;
  //   // assert( g_idx < _size);
  //   return _lst_bits[ g_idx] & ( 1LL << l_idx);
  // }

  void Set( _key_t idx, unsigned  int bit) {
    size_t g_idx = idx >> 5;
    int l_idx = idx & 0x1f;
    // assert( g_idx < _size);
#pragma omp atomic
    _lst_bits[ g_idx] |= bit << l_idx;
  }

  // void Set( long long int idx, long long int bit) {
  //   size_t g_idx = ( idx >> 6) & 0xffffffffffffffffL;
  //   long long int l_idx = idx & 0x3f;
  //   // assert( g_idx < _size);
  //   _lst_bits[ g_idx] |= bit << l_idx;
  // }

  void ClearAll() {
    for ( size_t i = 0; i < _size; i++) {
      _lst_bits[ i] = 0;
    }
  }

  _key_t CountBits() {
    // from http://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetParallel
//     int chunk = ( int)_size >> 3;
//     if ( !chunk) chunk = ( int)_size;
// #pragma omp parallel for schedule(static,chunk) reduction(+: tot)
    if ( ( omp_get_max_threads() < 2) || ( _size < 1024 * 1024)) {
      _key_t tot = 0;
      for ( int i = 0; i < ( int)_size; i++) {
        unsigned int val = _lst_bits[ i];
        val = val - ( ( val >> 1) & 0x55555555);                      // reuse input as temporary
        val = ( val & 0x33333333) + ( ( val >> 2) & 0x33333333);      // temp
        tot += ( ( ( val + ( val >> 4)) & 0xf0f0f0f) * 0x1010101) >> 24; // count
      }
      return tot;
    } else {
      _key_t tot = 0;
      int num_threads = omp_get_max_threads();
      int chunk = ( int)_size / num_threads;
      if ( !chunk) chunk = ( int)_size;
#pragma omp parallel
      {
        unsigned int mi_total = 0;
#pragma omp for schedule(static,chunk)
        for ( int i = 0; i < ( int)_size; i++) {
          unsigned int val = _lst_bits[ i];
          val = val - ( ( val >> 1) & 0x55555555);                      // reuse input as temporary
          val = ( val & 0x33333333) + ( ( val >> 2) & 0x33333333);      // temp
          mi_total += ( ( ( val + ( val >> 4)) & 0xf0f0f0f) * 0x1010101) >> 24; // count
        }
#pragma omp atomic
        tot += mi_total;
      }
      return tot;
    }
  }

  _key_t *get_occupied_indices();
  
 private:
  size_t _size;
  unsigned int *_lst_bits;
}; // BitArray

inline void fill_from_entry( unsigned int lst[ 32], unsigned int v) {
  lst[  0] = ( v & 0x00000001);
  lst[  1] = ( v & 0x00000002) >>  1;
  lst[  2] = ( v & 0x00000004) >>  2;
  lst[  3] = ( v & 0x00000008) >>  3;
  lst[  4] = ( v & 0x00000010) >>  4;
  lst[  5] = ( v & 0x00000020) >>  5;
  lst[  6] = ( v & 0x00000040) >>  6;
  lst[  7] = ( v & 0x00000080) >>  7;
  lst[  8] = ( v & 0x00000100) >>  8;
  lst[  9] = ( v & 0x00000200) >>  9;
  lst[ 10] = ( v & 0x00000400) >> 10;
  lst[ 11] = ( v & 0x00000800) >> 11;
  lst[ 12] = ( v & 0x00001000) >> 12;
  lst[ 13] = ( v & 0x00002000) >> 13;
  lst[ 14] = ( v & 0x00004000) >> 14;
  lst[ 15] = ( v & 0x00008000) >> 15;
  lst[ 16] = ( v & 0x00010000) >> 16;
  lst[ 17] = ( v & 0x00020000) >> 17;
  lst[ 18] = ( v & 0x00040000) >> 18;
  lst[ 19] = ( v & 0x00080000) >> 19;
  lst[ 20] = ( v & 0x00100000) >> 20;
  lst[ 21] = ( v & 0x00200000) >> 21;
  lst[ 22] = ( v & 0x00400000) >> 22;
  lst[ 23] = ( v & 0x00800000) >> 23;
  lst[ 24] = ( v & 0x01000000) >> 24;
  lst[ 25] = ( v & 0x02000000) >> 25;
  lst[ 26] = ( v & 0x04000000) >> 26;
  lst[ 27] = ( v & 0x08000000) >> 27;
  lst[ 28] = ( v & 0x10000000) >> 28;
  lst[ 29] = ( v & 0x20000000) >> 29;
  lst[ 30] = ( v & 0x40000000) >> 30;
  lst[ 31] = ( v & 0x80000000) >> 31;
}

inline _key_t *BitArray::get_occupied_indices() {
  unsigned int tam_ret = ( unsigned int)CountBits();
  if ( !tam_ret) return NULL;

  if ( ( omp_get_max_threads() < 2) || ( _size < 1024 * 1024)) {
    _key_t *ret_lst = new _key_t[ tam_ret + 1];
    int ret_idx = 0;
    for ( unsigned int i = 0; i < _size; i++) {
      unsigned int local[ 32];
      fill_from_entry( local, _lst_bits[ i]);
      if ( _lst_bits[ i]) {
        for ( int j = 0; j < 32; j++) {
          ret_lst[ ret_idx] = ( _key_t)i * ( _key_t)32 + ( _key_t)j;
          ret_idx += local[ j];
        }
      }
    }
    return ret_lst;
  } else {
    _key_t *idx_lst[ 256];
    int tam_idx_lst[ 256];
    for ( int i = 0; i < 256; i++) {
      idx_lst[ i] = NULL;
      tam_idx_lst[ i] = 0;
    }
    int num_threads = omp_get_max_threads();
    for ( int i = 0; i < num_threads; i++) {
      idx_lst[ i] = new _key_t[ tam_ret + 1];
    }
    
#pragma omp parallel
    {
      int chunk = ( int)_size / num_threads;
      if ( chunk <= 0) chunk = ( int)_size;
      int mi_id =  omp_get_thread_num();
      int mi_tam_idx_lst = 0;
      _key_t *mi_idx_lst = idx_lst[ mi_id];
#pragma omp for schedule( static, chunk)
      for ( int i = 0; i < ( int)_size; i++) {
        unsigned int local[ 32];
        fill_from_entry( local, _lst_bits[ i]);
        if ( _lst_bits[ i]) {
          for ( int j = 0; j < 32; j++) {
            mi_idx_lst[ mi_tam_idx_lst] = ( _key_t)i * ( _key_t)32 + ( _key_t)j;
            mi_tam_idx_lst += local[ j];
          }
        }
      }
#pragma omp critical
      tam_idx_lst[ mi_id] = mi_tam_idx_lst;
    }
    unsigned int total_idx = tam_idx_lst[ 0];
    for ( int i = 1; i < num_threads; i++) {
      for ( int j = 0; j < tam_idx_lst[ i]; j++) {
        idx_lst[ 0][ total_idx + j] = idx_lst[ i][ j];
      }
      total_idx += tam_idx_lst[ i];
      delete[] idx_lst[ i];
    }
    assert( total_idx == tam_ret);
    return idx_lst[ 0];
  }
}

class TwoBitArray {
 public:
  TwoBitArray( _key_t tam): _size( 0), _lst_bits( NULL) {
    _size = ( tam + 31 ) >> 4;
    _lst_bits = new unsigned int[ _size];
  }
  TwoBitArray( const TwoBitArray &ba): _size( 0), _lst_bits( NULL) {
    _size = ba._size;
    if ( _size) {
      _lst_bits = new unsigned int[ _size];
      for ( size_t i = 0; i < _size; i++) {
        _lst_bits[ i] = ba._lst_bits[ i];
      }
    }
  }
  ~TwoBitArray() {
    delete[] _lst_bits;
  }

  TwoBitArray &operator=( const TwoBitArray &ba) {
    delete[] _lst_bits;
    _lst_bits = NULL;
    _size = ba._size;
    if ( _size) {
      _lst_bits = new unsigned int[ _size];
      for ( size_t i = 0; i < _size; i++) {
        _lst_bits[ i] = ba._lst_bits[ i];
      }
    }
    return *this;
  }

  _key_t size() const {
    return ( _key_t)_size << 4; // it's not exact, but...
  }

  unsigned int Get( _key_t idx) const {
    size_t g_idx = idx >> 4;
    int l_idx = ( int)( idx & 0x0f) << 1;
    // assert( g_idx < _size);
    return ( _lst_bits[ g_idx] & ( 3 << l_idx)) >> l_idx;
  }

  void Set( _key_t idx, unsigned  int two_bits) {
    size_t g_idx = idx >> 4;
    int l_idx = ( int)( idx & 0x0f) << 1;
    // assert( g_idx < _size);
    _lst_bits[ g_idx] |= two_bits << l_idx;
  }

  void Incr( _key_t idx) {
    size_t g_idx = idx >> 4;
    int l_idx = ( int)( idx & 0x0f) << 1;
    unsigned int val =  ( _lst_bits[ g_idx] & ( 3 << l_idx)) >> l_idx;
    unsigned int incr[ 4] = { 1, 2, 3, 3};
    unsigned int next = incr[ val];
    _lst_bits[ g_idx] |= next << l_idx;
  }

  void ClearAll() {
    for ( size_t i = 0; i < _size; i++) {
      _lst_bits[ i] = 0;
    }
  }

  // _key_t CountBits() {
  //   // from http://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetParallel
  //   _key_t tot = 0;
  //   for ( size_t i = 0; i < _size; i++) {
  //     unsigned int val = _lst_bits[ i];
  //     val = val - ( ( val >> 1) & 0x55555555);                      // reuse input as temporary
  //     val = ( val & 0x33333333) + ( ( val >> 2) & 0x33333333);      // temp
  //     tot += ( ( ( val + ( val >> 4)) & 0xf0f0f0f) * 0x1010101) >> 24; // count
  //   }
  //   return tot;
  // }
  
 private:
  size_t _size;
  unsigned int *_lst_bits;
}; // TwoBitArray

#endif // _BIT_ARRAY_H_
